package inference

import "fmt"

// ModelVersion represents the YOLO model version
type ModelVersion string

const (
	ModelVersionYOLOv8  ModelVersion = "YOLOv8"
	ModelVersionYOLOv11 ModelVersion = "YOLOv11"
	ModelVersionUnknown ModelVersion = "Unknown"
)

// IsValid checks if the model version is valid
func (v ModelVersion) IsValid() bool {
	switch v {
	case ModelVersionYOLOv8, ModelVersionYOLOv11, ModelVersionUnknown:
		return true
	default:
		return false
	}
}

// String returns the string representation of the model version
func (v ModelVersion) String() string {
	return string(v)
}

// ParseModelVersion parses a string to ModelVersion
func ParseModelVersion(s string) ModelVersion {
	switch s {
	case "YOLOv8", "yolov8", "v8":
		return ModelVersionYOLOv8
	case "YOLOv11", "yolov11", "v11":
		return ModelVersionYOLOv11
	case "auto", "":
		return ModelVersionUnknown
	default:
		return ModelVersionUnknown
	}
}

// ModelInfo contains information about the loaded model
type ModelInfo struct {
	Version    ModelVersion `json:"version"`
	NumClasses int          `json:"num_classes"`
	InputSize  int          `json:"input_size"`
	ClassNames []string     `json:"class_names,omitempty"`
}

// String returns a formatted string representation of ModelInfo
func (m ModelInfo) String() string {
	return fmt.Sprintf("ModelInfo{Version: %s, NumClasses: %d, InputSize: %d}",
		m.Version, m.NumClasses, m.InputSize)
}

// Box represents a bounding box
type Box struct {
	X      int `json:"x"`
	Y      int `json:"y"`
	Width  int `json:"width"`
	Height int `json:"height"`
}

// DetectionResult represents a single detection result
type DetectionResult struct {
	ClassID    int     `json:"class_id"`
	ClassName  string  `json:"class_name"`
	Confidence float32 `json:"confidence"`
	Box        Box     `json:"box"`
}

// InferenceStats contains inference statistics
type InferenceStats struct {
	ProcessTimeMs float64 `json:"process_time_ms"`
	ModelVersion  string  `json:"model_version"`
	NumDetections int     `json:"num_detections"`
}

// DefaultCOCOClasses returns the default COCO 80 class names
func DefaultCOCOClasses() []string {
	return []string{
		"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
		"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
		"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
		"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
		"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
		"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
		"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
		"cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
		"scissors", "teddy bear", "hair drier", "toothbrush",
	}
}
